import pandas as pd
from pandas import DataFrame
from app_config import get_engine_ts
import numpy as np

pd.set_option('display.max_rows', None)  # 显示所有行
pd.set_option('display.max_columns', None)  # 显示所有列

# 设置显示宽度，防止换行
pd.set_option('display.width', None)  # 自动检测控制台宽度
pd.set_option('display.max_colwidth', None)  # 不截断列内容

engine = get_engine_ts()


class FundInfo(object):
    name: str = None
    code: str = None
    amount: float = None
    date: str = None


class StockInfo(object):
    code: str = None
    name: str = None
    amount: float = None
    msg_buy: str = None
    msg_sell: str = None


def fund_amount(fund_info: FundInfo = None) -> dict[str, DataFrame]:
    column_cons: str = "成分券代码Constituent Code"
    xls_cons: str = fund_info.code + 'cons.xls'
    cons: DataFrame = pd.read_excel(xls_cons, sheet_name=xls_cons, dtype={column_cons: str})
    cons = index_weight(fund_info.date, cons, column_cons)

    column_weight = "成份券代码Constituent Code"
    xls_weight = fund_info.code + 'closeweight.xls'
    weight = pd.read_excel(xls_weight, sheet_name=xls_weight, dtype={column_weight: str})

    set_buy = set(cons[column_cons]) - set(weight[column_weight])
    set_sell = set(weight[column_weight]) - set(cons[column_cons])

    buy: DataFrame = cons[cons[column_cons].isin(set_buy)]
    sell: DataFrame = weight[weight[column_weight].isin(set_sell)]

    sell['卖出金额（亿）'] = sell["权重(%)weight"] * fund_info.amount / 100
    buy['买入金额（亿）'] = buy["权重(%)weight"] * fund_info.amount / 100

    print(sell.columns.tolist())
    print(buy)
    print(sell)
    buy = buy[['指数名称 Index Name', '成分券代码Constituent Code', '成分券名称Constituent Name', '权重(%)weight', '买入金额（亿）']]
    sell = sell[['指数名称 Index Name', '成份券代码Constituent Code', '成份券名称Constituent Name', '权重(%)weight', '卖出金额（亿）']]

    buy.to_excel('bs/' + fund_info.code + '_' + fund_info.name + 'buy.xlsx')
    sell.to_excel('bs/' + fund_info.code + '_' + fund_info.name + 'sell.xlsx')
    sell_ = {"buy": buy, "sell": sell}
    return sell_


def index_weight(trade_date, df_index, head_key):
    pre_000300_index = df_index
    pre_000300_index['symbol'] = pre_000300_index[head_key]

    query = f""" select symbol, ts_code from stock_basic """
    stock_basic = pd.read_sql(query, get_engine_ts())

    pre_000300_index = pd.merge(pre_000300_index, stock_basic, on='symbol', how='left')
    pre_000300_index['con_code'] = pre_000300_index['ts_code']
    hs300_ts_code = pre_000300_index['con_code'].tolist()

    # 构建查询语句
    query = f"""
           SELECT ts_code,close,total_share,free_share FROM `daily_basic`
           WHERE ts_code IN ({','.join(f"'{code}'" for code in hs300_ts_code)})
           AND trade_date = '{trade_date}'
           """
    # 执行查询并将结果转换为DataFrame
    stock_share = pd.read_sql_query(query, get_engine_ts())
    hs300 = pd.merge(pre_000300_index, stock_share, left_on="con_code", right_on="ts_code", how='left')

    hs300['calc_percent'] = hs300['free_share'] / hs300['total_share'] * 100

    # 定义生成 b 列的规则
    def generate_b(value):
        if value <= 15:
            return np.ceil(value)  # 上调至最接近的整数
        elif 15 < value <= 20:
            return 20
        elif 20 < value <= 30:
            return 30
        elif 30 < value <= 40:
            return 40
        elif 40 < value <= 50:
            return 50
        elif 50 < value <= 60:
            return 60
        elif 60 < value <= 70:
            return 70
        elif 70 < value <= 80:
            return 80
        else:
            return 100

    hs300['adjust_percent'] = hs300['calc_percent'].apply(generate_b)
    hs300['adjust_share'] = hs300['adjust_percent'] / 100 * hs300['total_share']
    hs300['adjust_amount'] = hs300['adjust_share'] * hs300['close']

    amount_sum = hs300['adjust_amount'].sum()
    hs300['权重(%)weight'] = hs300['adjust_amount'] / amount_sum * 100
    return hs300


if __name__ == '__main__':

    params: list[FundInfo] = []

    excel = pd.read_excel("fund_amount.xlsx", sheet_name="Sheet1",
                          dtype={'total_amount': float, 'date': str, 'code': str})
    for index, row in excel.iterrows():
        code_ = row['code']
        name_ = row['name']
        amount_ = row['total_amount']
        date_ = row['date']
        row_: FundInfo = FundInfo()
        row_.name = name_
        row_.amount = amount_
        row_.date = date_
        row_.code = code_
        params.append(row_)

    returns: list[dict[str, DataFrame]] = []

    for param in params:
        re = fund_amount(param)
        returns.append(re)

    stock_dict: dict[str, StockInfo] = {}

    for re in returns:
        buy_df = re.get("buy")
        sell_df = re.get("sell")

        for index, row in buy_df.iterrows():
            buy_stock_code = row['成分券代码Constituent Code']
            buy_stock_name = row['成分券名称Constituent Name']
            buy_index_name = row['指数名称 Index Name']
            buy_weigh = row['权重(%)weight']
            buy_amount = row['买入金额（亿）']

            sto: StockInfo
            if stock_dict.__contains__(buy_stock_code):
                sto = stock_dict.get(buy_stock_code)
            else:
                sto = StockInfo()
                sto.msg_buy = ''
                sto.msg_sell = ''
                sto.amount = 0.0
                sto.code = buy_stock_code
                sto.name = buy_stock_name
                stock_dict.update({buy_stock_code: sto})

            sto.amount = sto.amount + buy_amount
            # sto.msg_buy = sto.msg_buy + "[" + buy_index_name + ": " + str(buy_amount) + "] "
            sto.msg_buy = f"{sto.msg_buy} [{buy_index_name}: {buy_amount:.1f}]"

        for index, row in sell_df.iterrows():
            buy_stock_code = row['成份券代码Constituent Code']
            buy_stock_name = row['成份券名称Constituent Name']
            buy_index_name = row['指数名称 Index Name']
            buy_weigh = row['权重(%)weight']
            buy_amount = row['卖出金额（亿）']

            sto: StockInfo
            if stock_dict.__contains__(buy_stock_code):
                sto = stock_dict.get(buy_stock_code)
            else:
                sto = StockInfo()
                sto.msg_buy = ''
                sto.msg_sell = ''
                sto.amount = 0.0
                sto.code = buy_stock_code
                sto.name = buy_stock_name
                stock_dict.update({buy_stock_code: sto})

            sto.amount = sto.amount - buy_amount
            # sto.msg_sell = sto.msg_sell + "[" + buy_index_name + ": " + str(buy_amount) + "] "
            sto.msg_sell = f"{sto.msg_sell} [{buy_index_name}: {buy_amount:.1f}]"

    list_data = [{"name": p.name, "code": p.code, "bs_amount": p.amount, "buy": p.msg_buy, "sell": p.msg_sell} for p in stock_dict.values()]
    frame = pd.DataFrame(data=list_data)

    # 构建查询语句2.426248201
    query = f"""
           select sb.symbol as code, sb.ts_code as ts_code from stock_basic sb
           """

    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    frame = pd.merge(frame, result_df, on='code', how='left')

    filter_ts_code = frame['ts_code'].tolist()
    strStartDate = '20250526'
    strEndDate = '20250530'
    # 构建查询语句
    query = f"""
        SELECT ts_code,trade_date, amount FROM `daily`
        WHERE ts_code IN ({','.join(f"'{code}'" for code in filter_ts_code)})
        AND trade_date >= '{strStartDate}'
        AND trade_date <= '{strEndDate}'
        """

    # 执行查询并将结果转换为DataFrame
    result_df = pd.read_sql_query(query, engine)
    result_df['amount'] = result_df['amount'] / 100000
    grouped_df = result_df.groupby('ts_code')['amount'].mean().reset_index()
    grouped_df.columns = ['ts_code', '日均成交额']
    merge = pd.merge(frame, grouped_df, on='ts_code', how='left')

    merge['冲击系数'] = merge['bs_amount'] / merge['日均成交额']
    merge = merge.sort_values("冲击系数", ascending=False).reset_index(drop=True)
    merge.to_excel("result3.xlsx")